package com.zzxz.practicaltrainingsys.component;

import com.alibaba.fastjson.JSON;
import com.zzxz.practicaltrainingsys.entity.ChatMessage;
import com.zzxz.practicaltrainingsys.entity.UserList;
import com.zzxz.practicaltrainingsys.service.ChatMessageService;
import com.zzxz.practicaltrainingsys.service.OllamaService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import javax.websocket.*;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;

@Component
@ServerEndpoint("/ai-chat/{sessionId}")
public class AIChatWebSocket {

    // 存储所有在线用户的会话
    private static final ConcurrentHashMap<String, Session> sessions = new ConcurrentHashMap<>();
    // 存储sessionId到用户名的映射（如果需要显示用户名而不是sessionId）
    private static final ConcurrentHashMap<String, String> sessionToUsername = new ConcurrentHashMap<>();

    private static OllamaService ollamaService;
    private static ChatMessageService chatMessageService;

    @Autowired
    public void setOllamaService(OllamaService ollamaService) {
        AIChatWebSocket.ollamaService = ollamaService;
    }

    @Autowired
    public void setChatMessageService(ChatMessageService chatMessageService) {
        AIChatWebSocket.chatMessageService = chatMessageService;
    }

    @OnOpen
    public void onOpen(Session session, @PathParam("sessionId") String sessionId) {
        // 检查用户是否已经在线，避免重复连接
        if (!sessions.containsKey(sessionId)) {
            sessions.put(sessionId, session);

            // 如果需要显示用户名，可以从sessionId中提取或通过其他方式获取
            // 这里假设sessionId就是用户标识，你可以根据实际情况修改
            sessionToUsername.put(sessionId, sessionId);

            System.out.println("AI聊天用户连接：" + sessionId);

            // 广播更新后的在线用户列表给所有用户
            broadcastUserList();

            // 在控制台显示当前在线用户
            showUserList();
        }

        // 发送历史消息
        sendHistoryMessages(session, sessionId);
    }

    @OnClose
    public void onClose(@PathParam("sessionId") String sessionId) {
        if (sessionId != null && sessions.containsKey(sessionId)) {
            sessions.remove(sessionId);
            sessionToUsername.remove(sessionId);

            System.out.println("AI聊天用户断开连接：" + sessionId);

            // 广播更新后的在线用户列表给所有用户
            broadcastUserList();

            // 在控制台显示当前在线用户
            showUserList();
        }
    }

    @OnMessage
    public void onMessage(String message, @PathParam("sessionId") String sessionId) {
        try {
            System.out.println("接收到客户端消息：" + message);

            WebSocketMessage wsMessage = JSON.parseObject(message, WebSocketMessage.class);
            String userMessage = wsMessage.getContent();

            // 保存用户消息
            if (chatMessageService != null) {
                chatMessageService.saveMessage(sessionId, "USER", userMessage);
            }

            // 回显用户消息
            WebSocketMessage userEcho = new WebSocketMessage("user_message", userMessage, sessionId);
            String echoJson = JSON.toJSONString(userEcho);
            System.out.println("发送回显消息给用户：" + echoJson);
            sendMessage(sessionId, echoJson);

            // 调用 AI 生成回复
            if (ollamaService != null) {
                String aiResponse = ollamaService.generateResponse(userMessage);
                System.out.println("AI 回复内容：" + aiResponse);

                if (chatMessageService != null) {
                    chatMessageService.saveMessage(sessionId, "AI", aiResponse);
                }

                WebSocketMessage aiMessage = new WebSocketMessage("ai_response", aiResponse, sessionId);
                String aiJson = JSON.toJSONString(aiMessage);
                System.out.println("发送 AI 回复消息：" + aiJson);
                sendMessage(sessionId, aiJson);
            }

        } catch (Exception e) {
            e.printStackTrace();
            WebSocketMessage errorMessage = new WebSocketMessage("error", "处理消息时出现错误", sessionId);
            sendMessage(sessionId, JSON.toJSONString(errorMessage));
        }
    }


    @OnError
    public void onError(Session session, Throwable error) {
        error.printStackTrace();
    }

    /**
     * 显示当前在线用户列表（控制台输出）
     */
    private void showUserList() {
        System.out.println("------ AI聊天室 当前在线用户 ------");
        for (String username : sessionToUsername.values()) {
            System.out.println(username);
        }
        System.out.println("在线用户总数：" + sessions.size());
        System.out.println("-----------------------------------");
    }

    /**
     * 创建用户列表JSON字符串
     */
    private String createUserListJson() {
        ArrayList<String> userList = new ArrayList<>(sessionToUsername.values());
        UserList userListObj = new UserList();
        userListObj.setUserlist(userList);
        userListObj.setType("userlist");
        return JSON.toJSONString(userListObj);
    }

    /**
     * 向所有在线用户广播用户列表更新
     */
    private void broadcastUserList() {
        String userListJson = createUserListJson();
        try {
            for (Session session : sessions.values()) {
                if (session.isOpen()) {
                    session.getBasicRemote().sendText(userListJson);
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    /**
     * 向指定用户发送消息
     */
    private void sendMessage(String sessionId, String message) {
        Session session = sessions.get(sessionId);
        if (session != null && session.isOpen()) {
            try {
                System.out.println("真正发送给 [" + sessionId + "] 的消息内容：" + message);
                session.getBasicRemote().sendText(message);
            } catch (Exception e) {
                e.printStackTrace();
            }
        } else {
            System.out.println("Session 不存在或已关闭，无法发送消息给：" + sessionId);
        }
    }

    /**
     * 发送历史消息给新连接的用户
     */
    private void sendHistoryMessages(Session session, String sessionId) {
        if (chatMessageService == null) return;

        try {
            List<ChatMessage> historyMessages = chatMessageService.getRecentMessages(sessionId, 20);

            // 倒序发送历史消息（从旧到新）
            for (int i = historyMessages.size() - 1; i >= 0; i--) {
                ChatMessage msg = historyMessages.get(i);
                String type = "USER".equalsIgnoreCase(msg.getSenderType()) ? "user_message" : "ai_response";


                WebSocketMessage wsMessage = new WebSocketMessage(type, msg.getContent(), sessionId);
                wsMessage.setTimestamp(msg.getCreateTime().toEpochSecond(java.time.ZoneOffset.of("+8")) * 1000);

                session.getBasicRemote().sendText(JSON.toJSONString(wsMessage));
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    /**
     * 获取当前在线用户数量
     */
    public static int getOnlineUserCount() {
        return sessions.size();
    }

    /**
     * 获取所有在线用户列表
     */
    public static List<String> getOnlineUsers() {
        return new ArrayList<>(sessionToUsername.values());
    }
}